package com.feidee.fd.sml.algorithm.feature

import com.feidee.fd.sml.algorithm.component.feature.{ChiSqSelectEncoder, ChiSqSelectEncoderParam}
import com.feidee.fd.sml.algorithm.util.{TestingDataGenerator, ToolClass}
import org.scalatest.FunSuite

/**
  * @Author: dongguosheng
  * @Date: 2019/3/21 16:24
  * @Review songhaicheng
  */
class ChiSqSelectEncoderSuite extends FunSuite {
  val paramStr: String =
    """
      |{
      |	"inputCol": "features",
      |	"numTopFeatures": 2,
      | "labelCol": "clicked",
      |	"outputCol": "selectedFeatures",
      |	"fdr": 0.05,
      | "fwe": 0.05,
      |	"fpr": 0.05,
      |	"percentile": 1.0
      |}
    """.stripMargin

  val chiSqSelector = new ChiSqSelectEncoder()
  val param: ChiSqSelectEncoderParam = chiSqSelector.parseParam(new ToolClass().encrypt(paramStr))

  test("ChiSqSelectorEncoder parameter construction test") {
    assert(param.input_pt == null && param.output_pt == null && "clicked".equals(param.labelCol) &&
      "features".equals(param.inputCol) && "selectedFeatures".equals(param.outputCol) && param.numTopFeatures > 0
      && param.fpr == 0.05
    )
  }

  test("ChiSqSelectorEncoder training function test") {
    val testingData = TestingDataGenerator.makeChiSqSelectorData()
    testingData.printSchema()
    testingData.show(false)
    val model = chiSqSelector.train(param, testingData)
    model.transform(testingData).show(false)
    // 当标签列字段类型为数字时，PipelineModel 有 3 个（字符串将会有 4 个）
    assert(model.stages.length == 3)
  }
}
